# Copyright (c) Alibaba, Inc. and its affiliates.
from dataclasses import dataclass, field
from typing import List, Literal, Optional, Union

from datasets import enable_caching

from swift.llm import DATASET_MAPPING, register_dataset_info
from swift.utils import get_logger

logger = get_logger()


@dataclass
class DataArguments:
    """
    DataArguments 类是一个数据类，包含与数据集处理相关的各种参数。

    参数：
    dataset (List[str]): dataset_id、dataset_path 或 dataset_dir 的列表
    val_dataset (List[str]): 验证数据集的 dataset_id、dataset_path 或 dataset_dir 的列表
    split_dataset_ratio (float): 如果 val_dataset 为空，则拆分数据集进行验证的比例。默认值为 0.01。
    data_seed (Optional[int]): 数据集 shuffling 的种子。默认值为 None。
    dataset_num_proc (int): 用于数据加载和预处理的进程数。默认值为 1。
    stream (bool): 启用数据集流式传输的标志。默认值为 False。
    download_mode (Literal): 下载数据集的模式。默认值为“reuse_dataset_if_exists”。
    columns: 用于手动映射数据集的列。
    model_name (List[str]): 包含模型的中英文名称的列表模型。默认为 [None, None]。
    model_author (List[str]): 模型作者的中英文姓名列表。默认为 [None, None]。
    custom_dataset_info (Optional[str]): 自定义 dataset_info.json 文件的路径。默认为 None。
    """
    # 定义训练数据集的路径、ID 或文件夹名，支持多个数据源
    dataset: List[str] = field(default_factory=list,
                               metadata={'help': f'dataset choices: {list(DATASET_MAPPING.keys())}'})

    # 定义验证数据集，同样支持多个路径、ID 或目录名
    val_dataset: List[str] = field(default_factory=list,
                                   metadata={'help': f'dataset choices: {list(DATASET_MAPPING.keys())}'})

    # 如果未提供验证集，则从训练集中按此比例划分验证集
    split_dataset_ratio: float = 0.01

    # 数据打乱使用的随机种子，None 表示使用默认或外部指定的随机种子
    data_seed: Optional[int] = None

    # 预处理和加载数据时使用的进程数，用于加速数据处理
    dataset_num_proc: int = 1

    # 是否从缓存文件加载数据，True 可提升加载速度
    load_from_cache_file: bool = True

    # 是否对训练集进行打乱
    dataset_shuffle: bool = True

    # 是否对验证集进行打乱
    val_dataset_shuffle: bool = False

    # 是否开启数据流式加载（适用于大数据集或远程数据）
    streaming: bool = False

    # 多数据集混合时的采样概率列表
    interleave_prob: Optional[List[float]] = None

    # 多数据源混合时，停止条件的策略
    stopping_strategy: Literal['first_exhausted', 'all_exhausted'] = 'first_exhausted'

    # 打乱数据时的缓冲区大小
    shuffle_buffer_size: int = 1000

    # 数据下载策略，'reuse_dataset_if_exists' 表示优先使用本地缓存
    download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists'

    # 用于列名映射的设置（dict 或 json 字符串），可手动指定数据集中哪些列是文本、标签等
    columns: Optional[Union[dict, str]] = None

    # 是否在列对齐失败时抛出异常
    strict: bool = False

    # 是否删除未被模型使用的数据列（节省内存）
    remove_unused_columns: bool = True

    # 模型名称，包含中文名和英文名，例如 ['小黄', 'Xiao Huang']
    model_name: List[str] = field(default_factory=lambda: [None, None],
                                  metadata={'help': "e.g. ['小黄', 'Xiao Huang']"})

    # 模型作者名，包含中文名和英文名，例如 ['魔搭', 'ModelScope']
    model_author: List[str] = field(default_factory=lambda: [None, None],
                                    metadata={'help': "e.g. ['魔搭', 'ModelScope']"})

    # 自定义 dataset_info.json 文件路径列表，用于扩展数据集信息
    custom_dataset_info: List[str] = field(default_factory=list)  # .json

    def _init_custom_dataset_info(self):
        """
        注册自定义的数据集描述信息文件（dataset_info.json）
        """
        # 如果提供的是字符串，转换成列表
        if isinstance(self.custom_dataset_info, str):
            self.custom_dataset_info = [self.custom_dataset_info]

        # 遍历每个 json 路径，并注册信息
        for path in self.custom_dataset_info:
            register_dataset_info(path)

    def __post_init__(self):
        # 如果未设置数据随机种子，则尝试使用外部 seed 属性
        if self.data_seed is None:
            self.data_seed = self.seed  # 需要外部注入 self.seed 属性

        # 解析列名映射配置
        self.columns = self.parse_to_dict(self.columns)

        # 如果手动设置了验证集，或使用流式加载，禁用自动划分验证集
        if len(self.val_dataset) > 0 or self.streaming:
            self.split_dataset_ratio = 0.0
            if len(self.val_dataset) > 0:
                msg = 'len(args.val_dataset) > 0'
            else:
                msg = 'args.streaming is True'
            logger.info(f'Because {msg}, setting split_dataset_ratio: {self.split_dataset_ratio}')

        # 注册自定义 dataset_info 文件
        self._init_custom_dataset_info()

    def get_dataset_kwargs(self):
        return {
            'seed': self.data_seed,
            'num_proc': self.dataset_num_proc,
            'load_from_cache_file': self.load_from_cache_file,
            'streaming': self.streaming,
            'interleave_prob': self.interleave_prob,
            'stopping_strategy': self.stopping_strategy,
            'shuffle_buffer_size': self.shuffle_buffer_size,
            'use_hf': self.use_hf,
            'hub_token': self.hub_token,
            'download_mode': self.download_mode,
            'columns': self.columns,
            'strict': self.strict,
            'model_name': self.model_name,
            'model_author': self.model_author,
            'remove_unused_columns': self.remove_unused_columns,
        }
